package com.tgy.tree;

import java.util.Comparator;
import java.util.Objects;

/**
 * @Author: tgy
 * @Date: 2020-09-29 16:33
 */
public class AVLTree<E> extends BinarySearchTree<E> {


    public AVLTree() {

        this(null);
    }

    public AVLTree(Comparator<E> comparator) {
        super(comparator);
    }


    private static class AVLNode<E> extends Node<E> {

        int height = 1;


        public AVLNode(E element, Node<E> parent) {
            super(element, parent);
        }

        public AVLNode<E> left(){

            return (AVLNode<E>)left;
        }

        public AVLNode<E> right() {

            return (AVLNode<E>)right;
        }

        public AVLNode<E> parent() {

            return (AVLNode<E>) parent;
        }


        public void addHeight() {

            height = Math.max(nodeHeight(left()),nodeHeight(right())) + 1;
        }

        public static <E> int nodeHeight(AVLNode<E> node) {

            return Objects.isNull(node) ? 0: node.height;
        }

        public boolean isBalance() {

            return Math.abs(nodeHeight(left()) - nodeHeight(right())) <= 1;
        }
    }

    @Override
    protected Node<E> createNode(E element, Node<E> parent) {

        return new AVLNode<>(element,parent);
    }

    @Override
    protected void addAdjust(Node<E> node) {

        AVLNode<E> avlNode = ((AVLNode<E>) node).parent();

        while (Objects.nonNull(avlNode) ) {


            if (avlNode.isBalance()) {
                // 平衡，就调整高度
                avlNode.addHeight();
            }else {

                break;
            }

            avlNode = avlNode.parent();
        }

        if (Objects.isNull(avlNode)) {

            return;
        }

        AVLNode<E> pNode = null;
//        判断旋转方式 RR LL RL LR
        if (avlNode.left().height > avlNode.right().height) {

//            L
            pNode = avlNode.left();

            if (pNode.left().height > pNode.right().height) {

//                L 右旋转

//                rotate(avlNode,true);

            }else {
//                R
                rotate(pNode,false);
            }

            rotate(avlNode,true);
        }else {
//            R
            pNode = avlNode.right();

            if (pNode.left().height > pNode.right().height){
//                L
                rotate(pNode,true);
//                rotate(avlNode,false);
            }else {
//                R
//                rotate(avlNode,false);
            }
            rotate(avlNode,false);
        }

//        调整高度

    }

    /**
     *
     * @param node
     * @param flag
     *    false 左旋转，
     *    true 右旋转
     */
    private void rotate(Node<E> node,boolean flag) {

        if (Objects.isNull(node)) {

            return;
        }

        Node<E> pNode = null;

        if (flag) {
//            右旋转

            pNode = node.left;
        }else {
//            左旋转
            pNode = node.right;
        }

//        设置parent

        pNode.parent = node.parent;

        if (Objects.isNull(node.parent)) {

            root = pNode;
        }else {

            if (node.isLeft()) {

                node.parent.left = pNode;
            }else {

                node.parent.right = pNode;
            }
        }
        node.parent = pNode;

        if (flag) {

            if (Objects.nonNull(pNode.left)) {

                pNode.left.parent = node;
            }

            pNode.left = node;

        }else {

            if (Objects.nonNull(pNode.right)) {

                pNode.right.parent = node;
            }

            pNode.right = node;
        }

    }
}
